package com.tool.engine.export.write.allocator;

import com.tool.engine.export.mertics.Collector;
import com.tool.engine.export.write.WriteTask;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import java.util.Objects;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.function.Function;
import java.util.function.UnaryOperator;

@Slf4j
@Setter
public class DataWriteThread extends Thread {

    private final BlockingQueue<Pair<WriteTask, Function<Boolean, Boolean>>> queue = new ArrayBlockingQueue<>(1024);

    private Collector collector;

    public DataWriteThread(String name) {
        super(name);
    }

    @Override
    public void run() {
        for (; ; ) {
            boolean success = false;
            Pair<WriteTask, Function<Boolean, Boolean>> pair = null;
            try {
                pair = queue.take();
                if (pair.getKey() == null) {
                    log.warn("write task is null");
                    continue;
                }
                long writeStartTime = System.currentTimeMillis();
                pair.getKey().write();
                long writeEndTime = System.currentTimeMillis();
                log.debug("taskName:{} write time: {} ms size:{}", pair.getKey().getTaskName(), writeEndTime - writeStartTime, pair.getKey().getSize());
                collector.incrWriteTime(writeEndTime - writeStartTime, pair.getKey().getSize());
                success= true;
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                log.warn("{} was InterruptedException:{}", Thread.currentThread().getName(), e.getMessage());
            } catch (Exception e) {
                log.warn("{} was exception:{}", Thread.currentThread().getName(), e.getMessage());
            }finally {
                if (Objects.nonNull(pair) && Objects.nonNull(pair.getValue())) {
                    pair.getValue().apply(success);
                }
            }
        }
    }

    public void put(WriteTask task, UnaryOperator<Boolean> consumer) {
        try {
            Pair<WriteTask, Function<Boolean, Boolean>> pair = new Pair<>(task, consumer);
            queue.put(pair);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("{} was put task occur InterruptedException:{}", Thread.currentThread().getName(), e.getMessage());
        }
    }
}
